# -*- coding: utf-8 -*-
import os
import traceback
import numpy as np
import win32gui
from ctypes import windll
import win32ui
import win32con
from PIL import Image
import time
import threading
from functools import lru_cache, wraps
from concurrent.futures import ThreadPoolExecutor

# 全局缓存锁
_cache_lock = threading.Lock()
# 资源回收间隔（秒）
CLEANUP_INTERVAL = 300


def gdi_resource_manager(func):
    """GDI资源自动管理装饰器"""

    @wraps(func)
    def wrapper(hwnd, *args, **kwargs):
        try:
            return func(hwnd, *args, **kwargs)
        finally:
            # 定期清理无效资源
            if not hasattr(threading.current_thread(), 'last_cleanup'):
                threading.current_thread().last_cleanup = time.time()

            now = time.time()
            if now - threading.current_thread().last_cleanup > CLEANUP_INTERVAL:
                cleanup_resources()
                threading.current_thread().last_cleanup = now

    return wrapper


def cleanup_resources():
    """清理无效窗口资源"""
    current_thread = threading.current_thread()
    if hasattr(current_thread, 'gdi_pool'):
        invalid_hwnds = [hwnd for hwnd in current_thread.gdi_pool
                         if not win32gui.IsWindow(hwnd)]
        for hwnd in invalid_hwnds:
            del current_thread.gdi_pool[hwnd]


def get_thread_resources():
    """获取线程级GDI资源池"""
    current_thread = threading.current_thread()
    if not hasattr(current_thread, 'gdi_pool'):
        current_thread.gdi_pool = {}
    return current_thread.gdi_pool


@lru_cache(maxsize=100)
def get_window_rect_cached(hwnd):
    """带缓存的窗口区域获取（线程安全）"""
    try:
        return win32gui.GetWindowRect(hwnd)
    except:
        return (0, 0, 0, 0)


@lru_cache(maxsize=100)
def get_client_rect_cached(hwnd):
    """带缓存的客户区获取（线程安全）"""
    try:
        return win32gui.GetClientRect(hwnd)
    except:
        return (0, 0, 0, 0)


def has_title_bar(hwnd):
    """判断窗口是否有标题栏（带异常处理）"""
    try:
        style = win32gui.GetWindowLong(hwnd, win32con.GWL_STYLE)
        return (style & win32con.WS_CAPTION) == win32con.WS_CAPTION
    except:
        return False


@gdi_resource_manager
def screenshot(hwnd, left=0, top=0, right=0, bottom=0, filename=None, is_top=False):
    """优化后的多线程安全截图函数"""
    try:
        # 缓存键生成
        current_time = time.time()
        window_text = win32gui.GetWindowText(hwnd) if win32gui.IsWindow(hwnd) else ""
        cache_key = (hwnd, left, top, right, bottom, is_top, window_text)

        # 检查时间窗口缓存
        with _cache_lock:
            thread_local = threading.current_thread()
            if hasattr(thread_local, 'last_shot') and \
                    (current_time - thread_local.last_time < 0.05) and \
                    thread_local.last_key == cache_key:
                img = thread_local.last_shot
                if filename:
                    Image.fromarray(img[..., [2, 1, 0]]).save(filename, quality=90)
                return img.copy()

        # 主截图逻辑
        if not is_top:
            # 客户区截图模式
            rect = get_client_rect_cached(hwnd)
            width = rect[2] - rect[0]
            height = rect[3] - rect[1]

            # 资源复用逻辑
            resources = get_thread_resources().get(hwnd)
            if resources and (resources['width'], resources['height']) == (width, height):
                hwnd_dc, mfc_dc, save_dc, save_bitmap = (
                    resources['hwnd_dc'],
                    resources['mfc_dc'],
                    resources['save_dc'],
                    resources['save_bitmap']
                )
            else:
                # 创建新GDI资源
                hwnd_dc = win32gui.GetDC(hwnd)
                mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc)
                save_dc = mfc_dc.CreateCompatibleDC()
                save_bitmap = win32ui.CreateBitmap()
                save_bitmap.CreateCompatibleBitmap(mfc_dc, width, height)
                save_dc.SelectObject(save_bitmap)

                # 缓存资源
                get_thread_resources()[hwnd] = {
                    'hwnd_dc': hwnd_dc,
                    'mfc_dc': mfc_dc,
                    'save_dc': save_dc,
                    'save_bitmap': save_bitmap,
                    'width': width,
                    'height': height
                }

            # 执行截图
            if windll.user32.PrintWindow(hwnd, save_dc.GetSafeHdc(), 3) == 0:
                return None

            # 转换图像数据
            bmp_info = save_bitmap.GetInfo()
            bmp_str = save_bitmap.GetBitmapBits(True)
            img = np.frombuffer(bmp_str, dtype='uint8').reshape(
                (bmp_info['bmHeight'], bmp_info['bmWidth'], 4))
            img = img[top:bottom, left:right]
        else:
            # 窗口截图模式
            if has_title_bar(hwnd):
                win_rect = get_window_rect_cached(hwnd)
                cli_rect = get_client_rect_cached(hwnd)
                title_w = (win_rect[2] - win_rect[0]) - (cli_rect[2] - cli_rect[0])
                title_h = (win_rect[3] - win_rect[1]) - (cli_rect[3] - cli_rect[1])
                abs_left = left + win_rect[0] + title_w - 8
                abs_top = top + win_rect[1] + title_h - 8
            else:
                win_rect = get_window_rect_cached(hwnd)
                abs_left = left + win_rect[0]
                abs_top = top + win_rect[1]

            width, height = right - left, bottom - top

            # 桌面DC获取
            hdesktop = win32gui.GetDesktopWindow()
            desktop_dc = win32gui.GetWindowDC(hdesktop)
            mem_dc = win32ui.CreateDCFromHandle(desktop_dc).CreateCompatibleDC()

            # 创建临时位图
            __screenshot = win32ui.CreateBitmap()
            __screenshot.CreateCompatibleBitmap(win32ui.CreateDCFromHandle(desktop_dc), width, height)
            mem_dc.SelectObject(__screenshot)

            # 执行截图
            mem_dc.BitBlt((0, 0), (width, height),
                          win32ui.CreateDCFromHandle(desktop_dc),
                          (abs_left, abs_top), win32con.SRCCOPY)

            # 转换图像数据
            bmp_info = __screenshot.GetInfo()
            bmp_str = __screenshot.GetBitmapBits(True)
            img = np.frombuffer(bmp_str, dtype='uint8').reshape(
                (bmp_info['bmHeight'], bmp_info['bmWidth'], 4))

            # 释放临时资源
            mem_dc.DeleteDC()
            win32gui.DeleteObject(__screenshot.GetHandle())
            win32gui.ReleaseDC(hdesktop, desktop_dc)

        # 文件保存
        if filename:
            os.makedirs(os.path.dirname(filename), exist_ok=True)
            Image.fromarray(img[..., [2, 1, 0]]).save(filename, quality=90)

        # 更新线程缓存
        with _cache_lock:
            thread_local = threading.current_thread()
            thread_local.last_shot = img.copy()
            thread_local.last_time = current_time
            thread_local.last_key = cache_key

        return img

    except Exception as e:
        #print(f"Screenshot error: {str(e)}")
        return None
    finally:
        # 显式释放桌面DC资源
        if is_top and 'desktop_dc' in locals():
            win32gui.ReleaseDC(win32gui.GetDesktopWindow(), desktop_dc)


# 使用示例
if __name__ == "__main__":
    def test_screenshot():
        hwnd = win32gui.GetDesktopWindow()
        img = screenshot(hwnd, is_top=True, right=800, bottom=600)
        if img is not False:
            print(f"Screenshot success! Image shape: {img.shape}")


    with ThreadPoolExecutor(max_workers=3) as executor:
        futures = [executor.submit(test_screenshot) for _ in range(5)]
        [f.result() for f in futures]
